{
 "cells": [
  {
   "cell_type": "code",
   "id": "initial_id",
   "metadata": {
    "collapsed": true,
    "ExecuteTime": {
     "end_time": "2025-02-22T12:09:08.033345Z",
     "start_time": "2025-02-22T12:09:08.015626Z"
    }
   },
   "source": [
    "import torch\n",
    "from torchtext.data.utils import get_tokenizer\n",
    "from torchtext.vocab import build_vocab_from_iterator\n",
    "from torchtext.transforms import VocabTransform\n",
    "from torchtext.datasets import WikiText2\n",
    "from torch.utils.data import DataLoader\n"
   ],
   "execution_count": 8,
   "outputs": []
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-22T12:44:00.961052Z",
     "start_time": "2025-02-22T12:44:00.390806Z"
    }
   },
   "cell_type": "code",
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
    "# 加载数据集 root的默认值为：C:\\Users\\Administrator\\.cache\\torch\\text\n",
    "# 还要在后面拼接上datasets/[数据集名称]/\n",
    "train_data, val_data, test_data = WikiText2(root='./')\n",
    "\n",
    "# 创建 DataLoader\n",
    "train_loader = DataLoader(list(train_data), batch_size=32, shuffle=True)"
   ],
   "id": "8d3954155e36aa4b",
   "execution_count": 31,
   "outputs": []
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-22T12:29:20.793678Z",
     "start_time": "2025-02-22T12:29:19.645597Z"
    }
   },
   "cell_type": "code",
   "source": [
    "def tokenizer(text):\n",
    "    return text.split()\n",
    "\n",
    "def yield_tokens(data_iter):\n",
    "    for text in data_iter:\n",
    "        yield tokenizer(text)\n",
    "# 构建词汇表\n",
    "vocab = build_vocab_from_iterator(yield_tokens(train_data), specials=[\"<unk>\", \"<pad>\", \"<bos>\", \"<eos>\"], max_tokens=20000)\n",
    "# 设置默认索引为 <unk>\n",
    "vocab.set_default_index(vocab[\"<unk>\"])\n",
    "# 定义文本转换函数\n",
    "text_transform = VocabTransform(vocab)\n",
    "\n",
    "# 打印词汇表大小\n",
    "print(f\"Vocabulary size: {len(vocab)}\")"
   ],
   "id": "afc9d25bf866999e",
   "execution_count": 20,
   "outputs": []
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-22T12:30:54.014284Z",
     "start_time": "2025-02-22T12:30:53.997286Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 数据加载器函数\n",
    "def collate_batch(batch):\n",
    "    labels, texts = zip(*batch)\n",
    "    texts = [text_transform(tokenizer(text)) for text in texts]\n",
    "    labels = [1 if label == \"pos\" else 0 for label in labels]\n",
    "    return torch.tensor(labels, dtype=torch.float32).to(device), torch.tensor(texts, dtype=torch.int64).to(device)"
   ],
   "id": "c80c25dd1115ed0e",
   "execution_count": 23,
   "outputs": []
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-22T12:41:04.740489Z",
     "start_time": "2025-02-22T12:41:04.026623Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# 创建 DataLoader\n",
    "train_loader = DataLoader(list(train_data), batch_size=8)\n",
    "test_loader = DataLoader(list(test_data), batch_size=8)\n",
    "\n",
    "# print(train_loader)\n",
    "# 测试 DataLoader\n",
    "for texts in train_loader:\n",
    "    # print(f\"Labels: {labels}\")\n",
    "    print(f\"Texts: {texts}\")\n",
    "    break\n"
   ],
   "id": "3b3d720e6b934936",
   "execution_count": 29,
   "outputs": []
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-02-22T13:21:36.785423Z",
     "start_time": "2025-02-22T13:21:36.727708Z"
    }
   },
   "cell_type": "code",
   "source": [
    "\n",
    "import torch\n",
    "from torchtext.vocab import vocab\n",
    "from torchtext.transforms import VocabTransform\n",
    "from collections import OrderedDict\n",
    "vocab_obj = vocab(['a','b','c','d'])\n",
    "vocab_transform = VocabTransform(vocab_obj)\n",
    "output = vocab_transform([['a','b'],['a','b','c']])\n",
    "jit_vocab_transform = torch.jit.script(vocab_transform)\n",
    "\n",
    "output"
   ],
   "id": "116d1c2d6745f2e7",
   "execution_count": 35,
   "outputs": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
